//////////////////////////////////////////////////////////////////////
// Haar2DCUDA_kernel.cu
//
// Fangyang SHEN, VRLab, Beihang University
// me@shenfy.com
//
// (C) Copyright VRLab, Beihang University 2010.
//////////////////////////////////////////////////////////////////////

//////////////////////////////////////////////////////////////////////
//  Brief Description
//
//	CUDA implementation of 2D Haar wavelet encoding/decoding,
//	device code file.
//
//////////////////////////////////////////////////////////////////////
#include <gpgpu_stdafx.h>
#include <Haar2DCUDA.h>
#include <cuda.h>
#include <Logger.h>
#include <ImgTmpl.h>
#include <iostream>

#pragma comment (lib, "cudart.lib")
#pragma comment (lib, W32LIB_BIN_FILE)

using namespace grflib::image;

#ifndef INV_SQRT_2
#define INV_SQRT_2 0.70710678118654752440f;
#endif


//Device code

__global__ void Haar2DEncH(float *out, float *in, int nCh, int stride, int halfLen)
{
	int y = blockIdx.x;
	int x = threadIdx.x;
	int c = blockIdx.y;

	int i1, i2, i3, i4;

	i1 = y * stride + x * nCh + c;
	i2 = y * stride + (2 * x) * nCh + c;
	i3 = i2 + nCh; //y * stride + (2 * x + 1) * nCh + c;
	i4 = i1 + halfLen * nCh; //y * stride + (x + halfLen) * nCh + c;

	float d1, d2;
	d1 = in[i2];
	d2 = in[i3];

	out[i1] = (d1 + d2) * INV_SQRT_2;
	out[i4] = (d1 - d2) * INV_SQRT_2;
}

__global__ void Haar2DEncV(float *out, float *in, int nCh, int stride, int halfLen)
{
	int y = blockIdx.x;
	//int x = threadIdx.x;
	//int c = blockIdx.y;
	int hOffset = threadIdx.x * nCh + blockIdx.y;

	int i1, i2, i3, i4;

	i1 = y * stride + hOffset;
	i2 = (2 * y) * stride + hOffset;
	i3 = (2 * y + 1) * stride + hOffset;
	i4 = (y + halfLen) * stride + hOffset;

	float d1, d2;
	d1 = in[i2];
	d2 = in[i3];

	out[i1] = (d1 + d2) * INV_SQRT_2;
	out[i4] = (d1 - d2) * INV_SQRT_2;
}

__global__ void Haar2DDecV(float *out, float *in, int nCh, int stride, int halfLen)
{
	int y = blockIdx.x;
	//int x = threadIdx.x;
	//int c = blockIdx.y;
	int hOffset = threadIdx.x * nCh + blockIdx.y;

	int i1, i2, i3, i4;

	i1 = y * stride + hOffset;
	i2 = (2 * y) * stride + hOffset;
	i3 = (2 * y + 1) * stride + hOffset;
	i4 = (y + halfLen) * stride + hOffset;

	float d1, d2;
	d1 = in[i1];
	d2 = in[i4];

	out[i2] = (d1 + d2) * INV_SQRT_2;
	out[i3] = (d1 - d2) * INV_SQRT_2;
}

__global__ void Haar2DDecH(float *out, float *in, int nCh, int stride, int halfLen)
{
	int y = blockIdx.x;
	int x = threadIdx.x;
	int c = blockIdx.y;

	int i1, i2, i3, i4;

	i1 = y * stride + x * nCh + c;
	i2 = y * stride + (2 * x) * nCh + c;
	i3 = i2 + nCh; //y * stride + (2 * x + 1) * nCh + c;
	i4 = i1 + halfLen * nCh; //y * stride + (x + halfLen) * nCh + c;

	float d1, d2;
	d1 = in[i1];
	d2 = in[i4];

	out[i2] = (d1 + d2) * INV_SQRT_2;
	out[i3] = (d1 - d2) * INV_SQRT_2;
}

//Host code

void CHaar2DCUDA::Release()
{
	if (m_pdBuf1)
	{
		cudaFree(m_pdBuf1);
	}

	if (m_pdBuf2)
	{
		cudaFree(m_pdBuf2);
	}

	m_pdBuf1 = 0;
	m_pdBuf2 = 0;
}

bool CHaar2DCUDA::Init(UINT width, UINT height, UINT channel)
{
	Release();

	if (width != height)
		return false;

	m_width = width;
	m_channel = channel;

	cudaError_t err;
	UINT lineLen = w32lib::CImgTmpl<float>::BytesPerLine(m_width, m_channel);

	if (cudaSuccess != (err = cudaMalloc(&m_pdBuf1, lineLen * m_width /*height==width*/)))
	{
		TCHAR msg[200];
		w32lib::CharToWChar(msg, cudaGetErrorString(err));
		w32lib::CLogger::Instance()->Error(msg);
		return false;
	}

	if (cudaSuccess != (err = cudaMalloc(&m_pdBuf2, lineLen * m_width /*height==width*/)))
	{
		TCHAR msg[200];
		w32lib::CharToWChar(msg, cudaGetErrorString(err));
		w32lib::CLogger::Instance()->Error(msg);
		return false;
	}
	return true;
}

bool CHaar2DCUDA::CudaEncode(float *out, const float *in)
{
	if (!m_pdBuf1 || !m_pdBuf2)
	{
		w32lib::CLogger::Instance()->Error(_T("CUDA buffer not initialized!"));
		return false;
	}

	UINT lineLen = w32lib::CImgTmpl<float>::BytesPerLine(m_width, m_channel);
	cudaError_t err;

	//copy input to graphics card
	VCUDA( cudaMemcpy(m_pdBuf1, in, lineLen * m_width, cudaMemcpyHostToDevice) );

	//launch kernels
	dim3 grid(1,1,1), block(1,1,1);
	int curLength = m_width / 2;

	while (curLength >= 1)
	{
		grid.x = curLength * 2; grid.y = m_channel;
		block.x = curLength;
		Haar2DEncH<<<grid, block>>>(m_pdBuf2, m_pdBuf1, m_channel, lineLen / sizeof(float), curLength);

		grid.x = curLength; grid.y = m_channel;
		block.x = curLength * 2;
		Haar2DEncV<<<grid, block>>>(m_pdBuf1, m_pdBuf2, m_channel, lineLen / sizeof(float), curLength);

		curLength /= 2;
	}
	VCUDA( cudaGetLastError() );
	
	//copy output
	VCUDA( cudaMemcpy(out, m_pdBuf1, lineLen * m_width, cudaMemcpyDeviceToHost) );
	
	return true;
}

bool CHaar2DCUDA::CudaDecode(float *out, const float *in)
{
	if (!m_pdBuf1 || !m_pdBuf2)
	{
		w32lib::CLogger::Instance()->Error(_T("CUDA buffer not initialized!"));
		return false;
	}

	UINT lineLen = w32lib::CImgTmpl<float>::BytesPerLine(m_width, m_channel);
	cudaError_t err;

	//copy input to graphics card
	VCUDA( cudaMemcpy(m_pdBuf1, in, lineLen * m_width, cudaMemcpyHostToDevice) );

	//launch kernels
	dim3 grid(1,1,1), block(1,1,1);
	UINT curLength = 1;

	while (curLength <= m_width / 2)
	{
		grid.x = curLength; grid.y = m_channel;
		block.x = curLength * 2;
		Haar2DDecV<<<grid, block>>>(m_pdBuf2, m_pdBuf1, m_channel, lineLen / sizeof(float), curLength);

		grid.x = curLength * 2; grid.y = m_channel;
		block.x = curLength;
		Haar2DDecH<<<grid, block>>>(m_pdBuf1, m_pdBuf2, m_channel, lineLen / sizeof(float), curLength);
		
		curLength *= 2;
	}
	VCUDA( cudaGetLastError() );

	//copy output
	VCUDA( cudaMemcpy(out, m_pdBuf1, lineLen * m_width, cudaMemcpyDeviceToHost) );

	return true;
}